import numpy as np
import os
import pandas as pd
from paddle.vision.transforms import ColorJitter
from PIL import Image
from process_labels import encode_labels
# ai.baidu.com/forum/topic/show/953213
def image_aug():

    """
    增广第5类和第8类
    遍历训练集的csv文件，读取每一张图片和label，检查label中是否有第5类和第8类
    若有，则用paddle中的ColorJitter类进行数据增广，增广图片保存在ImageAug文件夹
    把增广的图片路径加入到训练集的csv文件中
    :return:
    """
    img_aug_path = '/home/aistudio/ImageAug'
    img_aug_list = []
    img_aug_label_list = []
    if not os.path.exists(img_aug_path):
        os.makedirs(img_aug_path, exist_ok=True)
    train_csv_path = '/home/aistudio/work/data_list'
    train＿data = pd.read_csv(os.path.join(os.getcwd(), "work", "data_list", "train.csv"), sep=' ')
    print('len(train＿data)', len(train＿data))
    train＿data = pd.DataFrame(train＿data)
    for _, row in train＿data.iterrows():
        train_img = Image.open(row[0])
        train_mask = Image.open(row[1])
        # 读取图片文件名
        img_name = row[0].split('/')[-1]
        img_name, image_suffix = img_name.split('.')
        label_image = encode_labels(train_mask)
        # 增广第5类和第8类
        if np.count_nonzero(label_image == 4) != 0 or np.count_nonzero(label_image == 7) != 0:
            aug_image = ColorJitter(0.4, 0.4, 0.4, 0.4)(train_img)
            image_file_path = os.path.join(img_aug_path, img_name + '1.' + image_suffix)
            aug_image.save(image_file_path)
            img_aug_list.append(image_file_path)
            img_aug_label_list.append(row[1])
            image_file_path = os.path.join(img_aug_path, img_name + '2.' + image_suffix)
            aug_image = ColorJitter(0.4, 0.4, 0.4, 0.4)(train_img)
            aug_image.save(image_file_path)
            img_aug_list.append(image_file_path)
            img_aug_label_list.append(row[1])
    aug_pd = pd.DataFrame({'image':img_aug_list, 'label':img_aug_label_list})
    concated_pd = pd.concat([aug_pd, train＿data])
    concated_pd.to_csv(os.path.join(train_csv_path, 'train_concat.csv'), sep=' ', header=True, index=False)
if __name__ == '__main__':
    image_aug()
